# 服务器，用来控制训练过程
import copy
import torch 
from trainers.fedbase.FedServerBase import FedServerBase
import logging

# logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logging.basicConfig()
logger = logging.getLogger()
logger.setLevel(logging.INFO)

class FedAmpServer(FedServerBase):
    # 配置学习过程中的一些参数。主要是联邦学习相关的参数。
    def __init__(self,worker_list,model,server_params,client_params):
        """初始化服务器。并设置相关的参数。
        因为测试与训练一体。需要将训练端链接到测试端之后。
        这样因为train_worker_list是变长的，可能涉及道一系列对后续的影响。
        Args:
            worker_list (tuple(train_worker_list,test_worker_list)): 包含训练客户端列表和测试客户端列表
            model (torch.model): 实现了torch.model类的子类
            server_params (dict): 服务器-联邦学习相关的参数
            client_params (dict): 客户端-一次训练相关的参数
        """  
        super(FedAmpServer, self).__init__(worker_list,model,server_params,client_params)

        # FedAmpServer算法控制参数--------------------------------------------------------------------------------------
        # 表示其他模型对本模型的影响比例.等于0相当于本地训练，不受其他服务器上的影响。
        if("ak" in server_params):
            self.ak = server_params["ak"]
        else:
            self.ak = 0.4
        # 表示模型差异波动对模型贡献度影响的变化幅度。值越大，变化幅度越大
        if("sigma" in server_params):
            self.sigma = server_params["sigma"]
        else:
            self.sigma = 5

        # 根据客户单创建云端的Uk。用来保存融合了其他客户端的模型权重矩阵
        self.message_matrix = []

        # FedAmpServer特殊修改参数--------------------------------------------------------------------------------------
        # fedper和fedamp不需要单独的客户端而是保留一个本地个性化的模型。直接在本地个性化模型上，利用本地测试集数据进行测试。
        # 所以需要将分隔号的train_worker_list和test_worker_list合并到一块，最大化利用好数据。
        self.train_worker_list = self.train_worker_list+self.test_worker_list
        # 并修改信息矩阵的值。fedamp无法计算信息增益
        # self.information_increase_matrix=torch.zeros([len(self.train_worker_list)])
        logger.info("FedAmpServer Init:ak={},sigma={}".format(self.ak,self.sigma))




    def aggregate_by_weight(self,model_state_dict_list):   
        """ 首先计算模型参数之间的相似度。
        然后根据模型之间的相似度，进行聚合（融合）\n
        将其他客户端的知识融入到当前的计算当中\n
        信息公式求导的结果作为权重。主要进行归一化。

        Args:
            model_state_dict_list (dict[tensor]): pytorch格式的模型参数

        Returns:
            message_matrix: 信息矩阵U_k
        """        
        size = len(model_state_dict_list)

        # 将矩阵展开成一维的矩阵放到链表d当中
        # 放到链表当中m
        param_list = []
        for model_state_dict in model_state_dict_list:
            params = []
            for name in model_state_dict:
                params.append(torch.flatten(model_state_dict[name]))
            param_list.append(torch.cat(params))
        
        # 计算权重矩阵。注意，这里的权重是，A`的权重，A的导数。不是1-A
        weight_matrix = torch.zeros((size,size))
        for i in range(size):
            for j in range(size):
                if i==j:
                    continue
                weight_matrix[i][j] = torch.exp(-torch.norm(param_list[i]-param_list[j])/self.sigma)

        # 对权重矩阵进行归一化。并添加ak的影响
        weight_sum = torch.sum(weight_matrix,dim=1)
        for i in range(size):
            weight_matrix[i] = self.ak * weight_matrix[i]/weight_sum[i]
        
        for i in range(size):
            weight_matrix[i][i]=1-self.ak 

        # 利用权重对模型进行聚合。对这个list进行修改
        self.message_matrix = copy.deepcopy(model_state_dict_list)
        # 将信息矩阵中的内容都清空为零
        for model_state_dict in self.message_matrix:
            for name in model_state_dict:
                model_state_dict[name].zero_()
        # 根据权重矩阵
        for i in range(size):
            for j in range(size):
                for name in self.message_matrix[i]:
                    self.message_matrix[i][name]+=weight_matrix[i][j]*model_state_dict_list[j][name]

        return self.message_matrix

   

    # 
    def start_federated(self):
        """定义了训练的整个过程。fedAmp与其他算法的流程不同。
        1. 获取本地参数
        2. 聚合本地参数
        3. 发送聚合后的参数
        4. 执行本地训练
        5. 执行测试。打印并保存结果
        """        
        # 客户端进行训练 迭代以下过程
        for round in range(self.rounds):
            results_model = []
            # results_optim = []
            running_accuracy = []

            # 返回训练参数。得到各个模型的值
            for worker in self.train_worker_list:
                results_model.append(worker.get_model_state_dict())
            
            # 根据当前的权重，计算服务器上U的值。self.attention_matrix 
            self.aggregate_by_weight(results_model)
            
            # 发送模型。为每一个客户端发送一组新的模型复制
            for worker,model_state_dict in zip(self.train_worker_list,self.message_matrix):
                worker.load_model_state_dict(model_state_dict)

            # 训练模型
            for worker in self.train_worker_list:
                 running_accuracy.append(worker.train())

            # 对模型进行验证。每个服务器上的数据集分成两份。
            # 训练数据集和测试数据集。
            result_test=0
            for worker in self.train_worker_list:
                result_test += worker.test()/len(self.train_worker_list)
            
            # 数据记录(计算信息矩阵、保存测试准确率、训练损失、训练准确率、当前的模型参数保存为旧模型的参数)
            # 本身不具有一个全局的统一模型，这种计算也是毫无意义的。
            # if round>0:
            #     self.calculate_imformation_increase(results_model,self.test_accuracy[-1],result_test)

            self.test_accuracy.append(result_test)
            running_state = torch.tensor(running_accuracy)
            running_state = torch.mean(running_state,dim=0)
            self.train_loss.append(running_state[0].item())
            self.train_accuracy.append(running_state[1].item())
            self.old_model_parameters = copy.deepcopy(list(self.model.parameters()))
            logger.info("FedAmpServer start_federated:round={},test_acc={:.2f},train_loss={:.2f},train_acc={:.2f}".format(round, result_test,running_state[0].item(),running_state[1].item()))